// Copyright (C) 2010 and onwards Noam Shazeer and Georges Harik
// Pieces of this code may be licenced under the Apache License, version 2.0 
//   from Harik Shazeer Labs
#include "restaurant.h"

ThetaDist * G_THETAS = 0;

void InitRestaurant() {
  if (G_THETAS) return;
  G_THETAS = new ThetaDist;
  G_THETAS->Init(3, 10, false);
}

void ThetaDist::Init(const vector<uint64>& thetas, 
		     const vector<double>& priors) {
  thetas_ = thetas;
  priors_ = priors;
  Normalize(priors_);
  for (uint64 i=0; i<thetas.size(); i++) {
    ln_priors_.push_back(log(priors_[i]));
    ln_theta_.push_back(log(thetas_[i]));
    ln_gamma_theta_.push_back(FastLnFactorial(thetas_[i] - 1));
  }
  delete [] factor_cache_;
  factor_cache_ = new NLCache(1 << 20);
}

void ThetaDist::Init(uint64 levels, uint64 base, bool use_zero) {
  vector<uint64> thetas;
  vector<double> priors;
  if (use_zero) {
    thetas.push_back(0); 
    priors.push_back(1.0);
  }
  uint64 current = 1;
  for (uint64 i=0; i<levels; i++) {
    thetas.push_back(current);
    priors.push_back(1.0);
    current *= base;
  }
  Init(thetas, priors);
}

// returns the ln of a factor in the probability of a partition under
// the generalized chinese restaurant process. 
// in particular 
//    ln[ GAMMA(theta) * theta ^ num_tables / GAMMA(theta + total_count) ] 
// The other factor in the partition probability is the product of the
//   gamma functions of all the table sizes. 
double RestaurantFactor(uint64 num_tables, uint64 total_count, 
			uint64 theta, 
			double ln_theta, 
			double ln_gamma_theta) {
  if (theta == 0) {
    if (num_tables == 1) return -FastLnFactorial(total_count - 1);
    return NINF;
  }
  double ret = 0.0
    + ln_gamma_theta // + FastLnFactorial(theta - 1) 
    + num_tables * ln_theta
    - FastLnFactorial(theta + total_count - 1) 
    ; 
  return ret;
}

// Like RestaurantFactor, but for a generalized chinese restaurant process
// with this prior distribution over theta. 
int64 ThetaDist::Factor(uint64 num_tables, uint64 total_count) const {
  uint64 cache_key = Fingerprint(num_tables, total_count);
  uint64 result;
  if (factor_cache_->Lookup(cache_key, &result)) return result;
    
  double ret = NINF;
  for (uint64 i=0; i<thetas_.size(); i++) {
    double rf = 
      RestaurantFactor(num_tables, total_count, 
		       thetas_[i], ln_theta_[i], ln_gamma_theta_[i]);
    ret = LnSum(ret, rf + ln_priors_[i]);
  }
  int64 int_ret = (int64) (1e6*ret);
  factor_cache_->Store(cache_key, int_ret);
  return int_ret;
}

// returns a change in DL - 
int64 RestaurantUpdate(uint64 old_count, uint64 new_count, 
		       uint64 *total_count, uint64 *num_tables, 
		       const ThetaDist & thetas) {
  int64 dl_diff = 0;
  if (old_count > 0) dl_diff += MLnF(old_count - 1);
  if (new_count > 0) dl_diff -= MLnF(new_count - 1);
  dl_diff += thetas.Factor(*num_tables, *total_count);
  *total_count += new_count - old_count;
  if (old_count > 0) (*num_tables)--;
  if (new_count > 0) (*num_tables)++;
  dl_diff -= thetas.Factor(*num_tables, *total_count);
  return dl_diff;
}

int64 ImmutableRestaurantUpdate
(uint64 old_count, uint64 new_count, 
 uint64 old_total_count, uint64 old_num_tables,
 const ThetaDist & thetas) {
  return RestaurantUpdate
    (old_count, new_count, &old_total_count, &old_num_tables, thetas);
}

// Update a diner inside a Table<Thing> with many restaurants
BIND_HASH(kTotalCount); BIND_HASH(kNumTables);

int64 DinerUpdate
(Table<Thing> * h, uint64 diner, uint64 key, int64 diff, int64 dlkey) {
  uint64 count_hash = Fingerprint(diner, key);
  uint64 total_count_hash = Fingerprint(diner, kTotalCount);
  uint64 num_tables_hash = Fingerprint(diner, kNumTables);

  uint64 old_count = TableIntegerGet(h, count_hash);
  if (diff == 0) return old_count;
  uint64 new_count = old_count + diff;
  uint64 total_count = TableIntegerGet(h, total_count_hash);
  uint64 num_tables = TableIntegerGet(h, num_tables_hash);

  uint64 dl_diff = 
    RestaurantUpdate(old_count, new_count, &total_count, &num_tables, *G_THETAS);

  TableIntegerAdd(h, dlkey, dl_diff);
  TableIntegerSet(h, count_hash, new_count);
  TableIntegerSet(h, total_count_hash, total_count);
  TableIntegerSet(h, num_tables_hash, num_tables);
  return new_count;
}

int64 PartialRestaurantDL(uint64 count, uint64 total_count, uint64 num_tables,
			  const ThetaDist &thetas) {
  int64 ret = 0;
  if (count > 0) ret -= MLnF(count-1);
  ret -= thetas.Factor(num_tables, total_count);
  return ret;
}
